#!/usr/bin/env python3
"""
简单的WebSocket服务器示例
用于演示患者管理系统的实时通信功能
"""

import asyncio
import json
import logging
from datetime import datetime
from typing import Dict, Set

import websockets

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class PatientManagementServer:
    def __init__(self):
        self.clients: Set[websockets.WebSocketServerProtocol] = set()
        self.user_connections: Dict[str, websockets.WebSocketServerProtocol] = {}
        self.department_connections: Dict[str, Set[websockets.WebSocketServerProtocol]] = {}
        
    async def register_client(self, websocket: websockets.WebSocketServerProtocol):
        """注册新客户端"""
        self.clients.add(websocket)
        client_address = websocket.remote_address
        logger.info(f"新客户端连接: {client_address}，当前连接数: {len(self.clients)}")
        
    async def unregister_client(self, websocket: websockets.WebSocketServerProtocol):
        """注销客户端"""
        self.clients.discard(websocket)
        
        # 从用户连接映射中移除
        user_id = None
        for uid, conn in self.user_connections.items():
            if conn == websocket:
                user_id = uid
                break
        
        if user_id:
            del self.user_connections[user_id]
            
        # 从部门连接映射中移除（修复字典迭代错误）
        depts_to_remove = []
        for dept, connections in self.department_connections.items():
            connections.discard(websocket)
            if not connections:
                depts_to_remove.append(dept)
        
        # 安全地移除空部门
        for dept in depts_to_remove:
            del self.department_connections[dept]
                
        logger.info(f"客户端断开连接，当前连接数: {len(self.clients)}")
        
    async def handle_message(self, websocket: websockets.WebSocketServerProtocol, message: str):
        """处理接收到的消息"""
        try:
            data = json.loads(message)
            message_type = data.get('type')
            
            logger.info(f"收到消息类型: {message_type}")
            
            if message_type == 'AUTH':
                await self.handle_auth(websocket, data)
            elif message_type == 'PATIENT_REPORT':
                await self.handle_patient_report(websocket, data)
            elif message_type == 'NOTIFICATION':
                await self.handle_notification(websocket, data)
            elif message_type == 'HEARTBEAT':
                await self.handle_heartbeat(websocket, data)
            else:
                logger.warning(f"未知消息类型: {message_type}")
                # 发送错误响应
                error_response = {
                    'type': 'ERROR',
                    'data': {
                        'message': f'未知消息类型: {message_type}',
                        'timestamp': datetime.now().isoformat()
                    }
                }
                await websocket.send(json.dumps(error_response))
                
        except json.JSONDecodeError as e:
            logger.error(f"消息格式错误: {e}")
            error_response = {
                'type': 'ERROR',
                'data': {
                    'message': '消息格式错误',
                    'timestamp': datetime.now().isoformat()
                }
            }
            await websocket.send(json.dumps(error_response))
        except Exception as e:
            logger.error(f"处理消息时出错: {e}")
            error_response = {
                'type': 'ERROR',
                'data': {
                    'message': f'处理消息时出错: {str(e)}',
                    'timestamp': datetime.now().isoformat()
                }
            }
            await websocket.send(json.dumps(error_response))
            
    async def handle_auth(self, websocket: websockets.WebSocketServerProtocol, data: dict):
        """处理认证消息"""
        auth_data = data.get('data', {})
        user_id = str(auth_data.get('userId', ''))
        username = auth_data.get('username', '')
        user_type = auth_data.get('userType', '')
        department = auth_data.get('department', '')
        
        # 注册用户连接
        self.user_connections[user_id] = websocket
        
        # 注册部门连接
        if department:
            if department not in self.department_connections:
                self.department_connections[department] = set()
            self.department_connections[department].add(websocket)
            
        logger.info(f"用户认证成功: {username} ({user_type}) - {department}")
        
        # 发送认证确认
        response = {
            'type': 'AUTH_SUCCESS',
            'data': {
                'userId': user_id,
                'timestamp': datetime.now().isoformat()
            }
        }
        await websocket.send(json.dumps(response))
        
    async def handle_patient_report(self, websocket: websockets.WebSocketServerProtocol, data: dict):
        """处理患者上报消息"""
        report_data = data.get('data', {})
        patient_id = report_data.get('patientId', '')
        patient_name = report_data.get('patientName', '')
        reporter_name = report_data.get('reporterName', '')
        department = report_data.get('department', '')
        
        logger.info(f"收到患者上报: {patient_name} - {reporter_name} ({department})")
        
        # 广播给所有管理员和相关部门
        notification = {
            'type': 'PATIENT_REPORT',
            'data': {
                'patientId': patient_id,
                'patientName': patient_name,
                'reporterName': reporter_name,
                'department': department,
                'timestamp': datetime.now().isoformat()
            }
        }
        
        await self.broadcast_to_admins(notification)
        await self.broadcast_to_department(department, notification)
        
    async def handle_notification(self, websocket: websockets.WebSocketServerProtocol, data: dict):
        """处理通知消息"""
        notification_data = data.get('data', {})
        title = notification_data.get('title', '')
        content = notification_data.get('content', '')
        target_user_id = notification_data.get('targetUserId')
        target_department = notification_data.get('targetDepartment')
        
        logger.info(f"收到通知: {title} - {content}")
        
        notification = {
            'type': 'NOTIFICATION',
            'data': {
                'title': title,
                'content': content,
                'timestamp': datetime.now().isoformat()
            }
        }
        
        # 发送给指定用户
        if target_user_id:
            target_connection = self.user_connections.get(str(target_user_id))
            if target_connection:
                await target_connection.send(json.dumps(notification))
                
        # 发送给指定部门
        elif target_department:
            await self.broadcast_to_department(target_department, notification)
            
        # 广播给所有用户
        else:
            await self.broadcast_to_all(notification)
            
    async def handle_heartbeat(self, websocket: websockets.WebSocketServerProtocol, data: dict):
        """处理心跳消息"""
        response = {
            'type': 'HEARTBEAT_RESPONSE',
            'data': {
                'timestamp': datetime.now().isoformat()
            }
        }
        await websocket.send(json.dumps(response))
        
    async def broadcast_to_all(self, message: dict):
        """广播给所有客户端"""
        if not self.clients:
            return
            
        message_str = json.dumps(message)
        await asyncio.gather(
            *[client.send(message_str) for client in self.clients],
            return_exceptions=True
        )
        
    async def broadcast_to_admins(self, message: dict):
        """广播给管理员"""
        admin_connections = set()
        for user_id, connection in self.user_connections.items():
            # 这里可以根据用户类型判断是否为管理员
            # 简化处理，假设所有连接的用户都是管理员
            admin_connections.add(connection)
            
        if admin_connections:
            message_str = json.dumps(message)
            await asyncio.gather(
                *[client.send(message_str) for client in admin_connections],
                return_exceptions=True
            )
            
    async def broadcast_to_department(self, department: str, message: dict):
        """广播给指定部门"""
        if department not in self.department_connections:
            return
            
        connections = self.department_connections[department]
        if connections:
            message_str = json.dumps(message)
            await asyncio.gather(
                *[client.send(message_str) for client in connections],
                return_exceptions=True
            )
            
    async def handler(self, websocket: websockets.WebSocketServerProtocol, path: str):
        """WebSocket连接处理器"""
        await self.register_client(websocket)
        
        try:
            async for message in websocket:
                await self.handle_message(websocket, message)
        except websockets.exceptions.ConnectionClosed:
            logger.info("连接正常关闭")
        except Exception as e:
            logger.error(f"连接处理出错: {e}")
        finally:
            await self.unregister_client(websocket)
    
    async def handle_connection(self, websocket: websockets.WebSocketServerProtocol):
        """兼容性处理器"""
        await self.handler(websocket, "")

async def main():
    """主函数"""
    server = PatientManagementServer()
    
    # 启动WebSocket服务器
    host = "0.0.0.0"
    port = 8080
    logger.info(f"启动WebSocket服务器: ws://{host}:{port}")
    
    async with websockets.serve(server.handle_connection, host, port):
        logger.info("WebSocket服务器启动成功")
        await asyncio.Future()  # 保持服务器运行

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("服务器停止")
    except Exception as e:
        logger.error(f"服务器运行出错: {e}") 